{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial: Learning how to use Tune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"tune.png\" alt=\"Tune Logo\" width=\"400\"/>\n",
    "\n",
    "\n",
    "Tuning hyperparameters is often the most expensive part of the machine learning workflow. Tune is built to address this, demonstrating an efficient and scalable solution for this pain point. Note that this example depends on Tensorflow 2.0.\n",
    "\n",
    "**Code**: https://github.com/ray-project/ray/tree/master/python/ray/tune\n",
    "\n",
    "**Examples**: https://github.com/ray-project/ray/tree/master/python/ray/tune/examples\n",
    "\n",
    "**Documentation**: http://ray.readthedocs.io/en/latest/tune.html\n",
    "\n",
    "**Mailing List** https://groups.google.com/forum/#!forum/ray-dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you are running on Google Colab, uncomment below to install the necessary dependencies \n",
    "## before beginning the exercise.\n",
    "\n",
    "# print(\"Setting up colab environment\")\n",
    "# !pip uninstall -y -q pyarrow\n",
    "# !pip install -q -U ray[tune]\n",
    "# !pip install -q ray[debug]\n",
    "\n",
    "# # A hack to force the runtime to restart, needed to include the above dependencies.\n",
    "# print(\"Done installing! Restarting via forced crash (this is not an issue).\")\n",
    "# import os\n",
    "# os._exit(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you are running on Google Colab, please install TensorFlow 2.0 by uncommenting below..\n",
    "\n",
    "# try:\n",
    "#   # %tensorflow_version only exists in Colab.\n",
    "#   %tensorflow_version 2.x\n",
    "# except Exception:\n",
    "#   pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial will step through a couple key steps of the hyperparameter tuning process with Tune. \n",
    "\n",
    "1. Visualizing the data.\n",
    "2. Creating a model training procedure (using Keras).\n",
    "3. Tuning the model by adapting the above model training procedure to **use Tune**.\n",
    "4. Analyzing the model created by Tune.\n",
    "\n",
    "Note that this uses Tune's **function-based API**. This is mainly for prototyping. A later tutorial will cover Tune's more powerful **class-based Trainable** API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "\n",
    "import tensorflow as tf\n",
    "try:\n",
    "    tf.get_logger().setLevel('INFO')\n",
    "except Exception as exc:\n",
    "    print(exc)\n",
    "import warnings\n",
    "warnings.simplefilter(\"ignore\")\n",
    "\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense\n",
    "\n",
    "from tensorflow.keras.optimizers import SGD, Adam\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint\n",
    "\n",
    "import ray\n",
    "from ray import tune\n",
    "from ray.tune.examples.utils import get_iris_data\n",
    "\n",
    "import inspect\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('ggplot')\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize your data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first take a look at the distribution of the dataset.\n",
    "\n",
    "The Iris data sets consists of 3 different types of iris flowers’ (Setosa, Versicolour, and Virginica) petal and sepal length, stored in a 150x4 numpy.ndarray,\n",
    "\n",
    "The rows being the samples and the columns being: Sepal Length, Sepal Width, Petal Length and Petal Width.\n",
    "\n",
    "### The goal of this tutorial is to have a model that can accurately predict the true label given a 4-tuple of Sepal Length, Sepal Width, Petal Length and Petal Width."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets import load_iris\n",
    "\n",
    "iris = load_iris()\n",
    "true_data = iris['data']\n",
    "true_label = iris['target']\n",
    "names = iris['target_names']\n",
    "feature_names = iris['feature_names']\n",
    "\n",
    "def plot_data(X, y):\n",
    "    # Visualize the data sets\n",
    "    plt.figure(figsize=(16, 6))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    for target, target_name in enumerate(names):\n",
    "        X_plot = X[y == target]\n",
    "        plt.plot(X_plot[:, 0], X_plot[:, 1], linestyle='none', marker='o', label=target_name)\n",
    "    plt.xlabel(feature_names[0])\n",
    "    plt.ylabel(feature_names[1])\n",
    "    plt.axis('equal')\n",
    "    plt.legend();\n",
    "\n",
    "    plt.subplot(1, 2, 2)\n",
    "    for target, target_name in enumerate(names):\n",
    "        X_plot = X[y == target]\n",
    "        plt.plot(X_plot[:, 2], X_plot[:, 3], linestyle='none', marker='o', label=target_name)\n",
    "    plt.xlabel(feature_names[2])\n",
    "    plt.ylabel(feature_names[3])\n",
    "    plt.axis('equal')\n",
    "    plt.legend();\n",
    "    \n",
    "plot_data(true_data, true_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a model training procedure (using Keras)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's define a function that will take in some hyperparameters and return a model that we can then use to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model(learning_rate, dense_1, dense_2):\n",
    "    assert learning_rate > 0 and dense_1 > 0 and dense_2 > 0, \"Did you set the right configuration?\"\n",
    "    model = Sequential()\n",
    "    model.add(Dense(int(dense_1), input_shape=(4,), activation='relu', name='fc1'))\n",
    "    model.add(Dense(int(dense_2), activation='relu', name='fc2'))\n",
    "    model.add(Dense(3, activation='softmax', name='output'))\n",
    "    optimizer = SGD(lr=learning_rate)\n",
    "    model.compile(optimizer, loss='categorical_crossentropy', metrics=['accuracy'])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a function that trains the model using the ``create_model`` function and returns the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_on_iris():\n",
    "    train_x, train_y, test_x, test_y = get_iris_data()\n",
    "    model = create_model(learning_rate=0.1, dense_1=2, dense_2=2)\n",
    "    # This saves the top model. `accuracy` is only available in TF2.0.\n",
    "    checkpoint_callback = ModelCheckpoint(\n",
    "        \"model.h5\", monitor='accuracy', save_best_only=True, save_freq=2)\n",
    "\n",
    "    # Train the model\n",
    "    model.fit(\n",
    "        train_x, train_y, \n",
    "        validation_data=(test_x, test_y),\n",
    "        verbose=0, batch_size=10, epochs=20, callbacks=[checkpoint_callback])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's quickly train the model on the dataset. The accuracy should be quite low."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_model = train_on_iris()  # This trains the model and returns it.\n",
    "train_x, train_y, test_x, test_y = get_iris_data()\n",
    "original_loss, original_accuracy = original_model.evaluate(test_x, test_y, verbose=0)\n",
    "print(\"Loss is {:0.4f}\".format(original_loss))\n",
    "print(\"Accuracy is {:0.4f}\".format(original_accuracy))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Integrate with Tune\n",
    "\n",
    "Now, let's use Tune to optimize a model that learns to classify Iris. This will happen in two parts - **modifying** the training function to support Tune, and then **configuring** Tune.\n",
    "\n",
    "Let's first define a callback function to report intermediate training progress back to Tune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow.keras as keras\n",
    "from ray import tune\n",
    "\n",
    "\n",
    "class TuneReporterCallback(keras.callbacks.Callback):\n",
    "    \"\"\"Tune Callback for Keras.\n",
    "    \n",
    "    The callback is invoked every epoch.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, logs={}):\n",
    "        self.iteration = 0\n",
    "        super(TuneReporterCallback, self).__init__()\n",
    "\n",
    "    def on_epoch_end(self, batch, logs={}):\n",
    "        self.iteration += 1\n",
    "        tune.report(keras_info=logs, mean_accuracy=logs.get(\"accuracy\"), mean_loss=logs.get(\"loss\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integration Part 1: Modifying the Training Function\n",
    "\n",
    "**Instructions** Follow the next 2 steps for modifying the ``train_iris`` function to support Tune.\n",
    "\n",
    "1. Change the signature of the function to take in a hyperparameter dictionary. This function will be called on a Ray.\n",
    "\n",
    "```python\n",
    "def tune_iris(config)\n",
    "```\n",
    "    \n",
    "    \n",
    "2. Pass in the configuration values into ``create_model``:\n",
    "\n",
    "```python\n",
    "model = create_model(learning_rate=config[\"lr\"], dense_1=config[\"dense_1\"], dense_2=config[\"dense_2\"])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tune_iris():  # TODO: Change me.\n",
    "    train_x, train_y, test_x, test_y = get_iris_data()\n",
    "    model = create_model(learning_rate=0, dense_1=0, dense_2=0)  # TODO: Change me.\n",
    "    checkpoint_callback = ModelCheckpoint(\n",
    "        \"model.h5\", monitor='loss', save_best_only=True, save_freq=2)\n",
    "\n",
    "    # Enable Tune to make intermediate decisions by using a Tune Callback hook. This is Keras specific.\n",
    "    callbacks = [checkpoint_callback, TuneReporterCallback()]\n",
    "    \n",
    "    # Train the model\n",
    "    model.fit(\n",
    "        train_x, train_y, \n",
    "        validation_data=(test_x, test_y),\n",
    "        verbose=0, \n",
    "        batch_size=10, \n",
    "        epochs=20, \n",
    "        callbacks=callbacks)\n",
    "    \n",
    "assert len(inspect.getargspec(tune_iris).args) == 1, \"The `tune_iris` function needs to take in the arg `config`.\"\n",
    "\n",
    "# print(\"Test-running to make sure this function will run correctly.\")\n",
    "# tune.track.init()  # For testing purposes only.\n",
    "# tune_iris({\"lr\": 0.1, \"dense_1\": 4, \"dense_2\": 4})\n",
    "# print(\"Success!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integration Part 2: Configuring Tune to tune hyperparameters.\n",
    "\n",
    "**Instructions** Follow the next 2 steps to configure Tune to identify the top hyperparameters.\n",
    "\n",
    "1. Designate the hyperparameter space. \n",
    "\n",
    "```python\n",
    "hyperparameter_space = {\n",
    "    \"lr\": tune.loguniform(0.001, 0.1),  \n",
    "    \"dense_1\": tune.uniform(2, 128),\n",
    "    \"dense_2\": tune.uniform(2, 128),\n",
    "}\n",
    "```\n",
    "2. Increase the number of samples. The more trials we evaluate, the more chances we choose a good model.\n",
    "\n",
    "```python\n",
    "num_samples = 20\n",
    "```\n",
    "#### FAQ: How does parallelism work in Tune?\n",
    "\n",
    "\n",
    "Setting ``num_samples`` will run a *total* of 20 trials (hyperparameter configuration samples). However, not all of them will run at once. The max training concurrency will be the number of CPU cores on the machine you're running on. For a 2-core machine, 2 models will be trained concurrently. When one is finished, a new training process will start with a new hyperparameter configuration sample. \n",
    "\n",
    "Each trial will run on a new Python process. The python process is killed when the trial is finished.\n",
    "\n",
    "\n",
    "#### FAQ: How do I debug things in Tune?\n",
    "\n",
    "The `error file` column will show up in the output. Run the below cell with the ``error file path`` path to diagnose your issue.\n",
    "\n",
    "```\n",
    "! cat /home/ubuntu/tune_iris/tune_iris_c66e1100_2019-10-09_17-13-24x_swb9xs/error_2019-10-09_17-13-29.txt\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Launch a Tune hyperparameter search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# This seeds the hyperparameter sampling.\n",
    "import numpy as np; np.random.seed(5)  \n",
    "hyperparameter_space = {}  # TODO: Fill me out.\n",
    "num_samples = 1  # TODO: Fill me out.\n",
    "\n",
    "####################################################################################################\n",
    "################ This is just a validation function for tutorial purposes only. ####################\n",
    "HP_KEYS = [\"lr\", \"dense_1\", \"dense_2\"]\n",
    "assert all(key in hyperparameter_space for key in HP_KEYS), (\n",
    "    \"The hyperparameter space is not fully designated. It must include all of {}\".format(HP_KEYS))\n",
    "######################################################################################################\n",
    "\n",
    "ray.shutdown()  # Restart Ray defensively in case the ray connection is lost. \n",
    "ray.init(log_to_driver=False)\n",
    "# We clean out the logs before running for a clean visualization later.\n",
    "! rm -rf ~/ray_results/tune_iris\n",
    "\n",
    "analysis = tune.run(\n",
    "    tune_iris, \n",
    "    verbose=1, \n",
    "    config=hyperparameter_space,\n",
    "    num_samples=num_samples)\n",
    "\n",
    "assert len(analysis.trials) == 20, \"Did you set the correct number of samples?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze the best tuned model\n",
    "\n",
    "Let's compare the real labels with the classified labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "_, _, test_data, test_labels = get_iris_data()\n",
    "plot_data(test_data, test_labels.argmax(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obtain the directory where the best model is saved.\n",
    "print(\"You can use any of the following columns to get the best model: \\n{}.\".format(\n",
    "    [k for k in analysis.dataframe() if k.startswith(\"keras_info\")]))\n",
    "print(\"=\" * 10)\n",
    "logdir = analysis.get_best_logdir(\"keras_info/val_loss\", mode=\"min\")\n",
    "# We saved the model as `model.h5` in the logdir of the trial.\n",
    "from tensorflow.keras.models import load_model\n",
    "tuned_model = load_model(logdir + \"/model.h5\")\n",
    "\n",
    "tuned_loss, tuned_accuracy = tuned_model.evaluate(test_data, test_labels, verbose=0)\n",
    "print(\"Loss is {:0.4f}\".format(tuned_loss))\n",
    "print(\"Tuned accuracy is {:0.4f}\".format(tuned_accuracy))\n",
    "print(\"The original un-tuned model had an accuracy of {:0.4f}\".format(original_accuracy))\n",
    "predicted_label = tuned_model.predict(test_data)\n",
    "plot_data(test_data, predicted_label.argmax(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare the performance of the best model by visualizing predictions compared to the ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_comparison(X, y):\n",
    "    # Visualize the data sets\n",
    "    plt.figure(figsize=(16, 6))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    for target, target_name in enumerate([\"Incorrect\", \"Correct\"]):\n",
    "        X_plot = X[y == target]\n",
    "        plt.plot(X_plot[:, 0], X_plot[:, 1], linestyle='none', marker='o', label=target_name)\n",
    "    plt.xlabel(feature_names[0])\n",
    "    plt.ylabel(feature_names[1])\n",
    "    plt.axis('equal')\n",
    "    plt.legend();\n",
    "\n",
    "    plt.subplot(1, 2, 2)\n",
    "    for target, target_name in enumerate([\"Incorrect\", \"Correct\"]):\n",
    "        X_plot = X[y == target]\n",
    "        plt.plot(X_plot[:, 2], X_plot[:, 3], linestyle='none', marker='o', label=target_name)\n",
    "    plt.xlabel(feature_names[2])\n",
    "    plt.ylabel(feature_names[3])\n",
    "    plt.axis('equal')\n",
    "    plt.legend();\n",
    "    \n",
    "plot_comparison(test_data, test_labels.argmax(1) == predicted_label.argmax(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extra - use Tensorboard for results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use TensorBoard to view trial performances. If the graphs do not load, click `Toggle All Runs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%tensorboard --logdir ~/ray_results/tune_iris"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
